package com.star.spring.websocket.handler;

import com.alibaba.fastjson2.JSON;
import com.star.spring.websocket.model.Message;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class WebSocketHandler implements org.springframework.web.socket.WebSocketHandler {

    // 存储所有连接的会话
    private static final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String sessionId = session.getId();
        sessions.put(sessionId, session);
        log.info("WebSocket连接建立，sessionId: {}", sessionId);
        
        // 发送欢迎消息
        Message welcomeMessage = new Message();
        welcomeMessage.setType("WELCOME");
        welcomeMessage.setContent("欢迎连接到WebSocket服务器！");
        welcomeMessage.setTimestamp(System.currentTimeMillis());
        sendMessage(sessionId, welcomeMessage);
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        String sessionId = session.getId();
        String payload = message.getPayload().toString();
        log.info("收到来自sessionId: {} 的消息: {}", sessionId, payload);
        
        try {
            // 解析接收到的消息
            Message receivedMessage = JSON.parseObject(payload, Message.class);
            
            // 处理消息
            handleReceivedMessage(sessionId, receivedMessage);
        } catch (Exception e) {
            log.error("解析消息失败: {}", payload, e);
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        String sessionId = session.getId();
        log.error("WebSocket传输错误，sessionId: {}", sessionId, exception);
        sessions.remove(sessionId);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        String sessionId = session.getId();
        log.info("WebSocket连接关闭，sessionId: {}, status: {}", sessionId, closeStatus);
        sessions.remove(sessionId);
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 处理接收到的消息
     */
    private void handleReceivedMessage(String sessionId, Message message) {
        switch (message.getType()) {
            case "ECHO":
                // 回显消息
                Message echoMessage = new Message();
                echoMessage.setType("ECHO_RESPONSE");
                echoMessage.setContent("服务器回显: " + message.getContent());
                echoMessage.setTimestamp(System.currentTimeMillis());
                sendMessage(sessionId, echoMessage);
                break;
            case "BROADCAST":
                // 广播消息给所有连接
                Message broadcastMessage = new Message();
                broadcastMessage.setType("BROADCAST");
                broadcastMessage.setContent("广播消息: " + message.getContent());
                broadcastMessage.setTimestamp(System.currentTimeMillis());
                broadcastMessage(broadcastMessage);
                break;
            default:
                log.warn("未知的消息类型: {}", message.getType());
        }
    }

    /**
     * 发送消息给指定会话
     */
    public void sendMessage(String sessionId, Message message) {
        WebSocketSession session = sessions.get(sessionId);
        if (session != null && session.isOpen()) {
            try {
                String jsonMessage = JSON.toJSONString(message);
                session.sendMessage(new TextMessage(jsonMessage));
                log.info("发送消息给sessionId: {}, 消息: {}", sessionId, jsonMessage);
            } catch (IOException e) {
                log.error("发送消息失败，sessionId: {}", sessionId, e);
                sessions.remove(sessionId);
            }
        }
    }

    /**
     * 广播消息给所有连接
     */
    public void broadcastMessage(Message message) {
        String jsonMessage = JSON.toJSONString(message);
        sessions.values().removeIf(session -> {
            if (session.isOpen()) {
                try {
                    session.sendMessage(new TextMessage(jsonMessage));
                    return false;
                } catch (IOException e) {
                    log.error("广播消息失败", e);
                    return true;
                }
            }
            return true;
        });
        log.info("广播消息: {}", jsonMessage);
    }

    /**
     * 获取当前连接数
     */
    public int getConnectionCount() {
        return sessions.size();
    }

    /**
     * 获取所有会话ID
     */
    public java.util.Set<String> getAllSessionIds() {
        return sessions.keySet();
    }
} 